/*
 * memory-profile-ptrace: A program for measuring memory utilisation via
 *  ptrace on Linux-x86.
 *
 * Copyright (C) 2007 Andrew Miller
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 */
#include <sys/types.h>
#include <unistd.h>
#include <errno.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <asm/unistd.h>
#include <sys/ptrace.h>
#include <sys/mman.h>
#include <arpa/inet.h>

class Process;
Process* processTable[0x10000];

unsigned long gAllocatedBytes = 0;
unsigned long gPeakAllocatedBytes = 0;

struct mmap_arg_struct
{
  unsigned long addr;
  unsigned long len;
  unsigned long prot;
  unsigned long flags;
  unsigned long fd;
  unsigned long offset;
};

void
PrintStatistics()
{
  printf("Currently allocated: %lu bytes. Peak: %lu bytes.\n",
         gAllocatedBytes, gPeakAllocatedBytes);
}

void
ChangeGlobalAllocation(int bytes)
{
  gAllocatedBytes += bytes;
  if (gAllocatedBytes > gPeakAllocatedBytes)
    gPeakAllocatedBytes = gAllocatedBytes;
}

class Process
{
public:
  Process(pid_t aPid, bool aInSyscall = false)
    : mInSyscall(aInSyscall), mIsMain(false), mPid(aPid),
      mAllocatedBytes(0), mLastBreak(0x8000000),
      mForkHandlingState(FHS_NEVER_HANDLED)
  {
  }

  void
  waitReturned(int status)
  {
    if (!WIFSTOPPED(status))
    {
      processExited();
      return;
    }

    if (WSTOPSIG(status) == SIGSEGV)
    {
      user_regs_struct regs;
      ptrace(PTRACE_GETREGS, mPid, NULL, &regs);
      fprintf(stderr, "Got segment violation at %08lX\n",
             regs.eip);
      fflush(stderr);

      exit(0);
    }

    if (WSTOPSIG(status) != SIGTRAP)
    {
      ptrace(PTRACE_SYSCALL, mPid, NULL, WSTOPSIG(status));
      return;
    }

    traceSyscall();
    
    ptrace(PTRACE_SYSCALL, mPid, NULL, NULL);
    mInSyscall = !mInSyscall;
  }

  void
  isMainProcess()
  {
    mIsMain = true;

    ptrace(PTRACE_SYSCALL, mPid, NULL, NULL);
  }
private:
  bool mInSyscall;
  bool mIsMain;
  pid_t mPid;
  unsigned long mAllocatedBytes;
  unsigned long mLastBreak;

  void
  processExited()
  {
    if (mIsMain)
    {
      PrintStatistics();
      printf("Main process exited.\n");
      exit(0);
    }

    ChangeGlobalAllocation(-mAllocatedBytes);

    // TODO: free resources held by the process.
    processTable[mPid] = NULL;
    delete this;
  }

  enum
  {
    FHS_NEVER_HANDLED,
    FHS_NOT_HANDLING,
    FHS_WAITING_MMAP2,
    FHS_WAITING_FORK,
    FHS_WAITING_NOOP
  } mForkHandlingState;

  unsigned long eax, ebx, ecx, edx, esi, eip;
  unsigned long mForkHandlerAddress;

  void
  doForkHandling(struct user_regs_struct& regs)
  {
    if (mForkHandlingState == FHS_WAITING_MMAP2)
    {
      mForkHandlerAddress = regs.eax;
      const char* forkCode =
        "\xb8\x02\x00\x00\x00" // movl $2, %eax   /* fork */
        "\xcd\x80"             // int $0x80       /* syscall */
        "\xb8\x14\x00\x00\x00" // movl $20, %eax  /* getpid */
        "\xcd\x80"             // int $0x80       /* syscall */
        "\x89\xc3"             // mov %eax, %ebx  /* send to self */
        "\xb8\x25\x00\x00\x00" // movl $37, %eax  /* kill */
        "\xb9\x13\x00\x00\x00" // movl $19, %ecx  /* SIGSTOP */
        "\xcd\x80"             // int $0x80       /* syscall */
        ;
      copy_to_user(mForkHandlerAddress, forkCode, 28);

      regs.eip = mForkHandlerAddress;

      // Save registers...
      ptrace(PTRACE_SETREGS, mPid, NULL, &regs);

      mForkHandlingState = FHS_WAITING_FORK;

      // Return into the fork handler...
      return;
    }

    if (!mInSyscall)
      return;

    if (mForkHandlingState == FHS_WAITING_NOOP)
    {
      mForkHandlingState = FHS_WAITING_FORK;
      return;
    }

    mForkHandlingState = FHS_NOT_HANDLING;

    // Okay, we have just returned from fork(). The new process is stopped, its
    // pid is in regs.eax.

    pid_t newPid = static_cast<pid_t>(regs.eax);

    if (newPid < 0)
    {
      fprintf(stderr, "Could fork child: %s\n", strerror(-newPid));
      fflush(stderr);
      return;
    }

    // Attach to the new process...
    ptrace(PTRACE_ATTACH, newPid, NULL, NULL);
    int status;
    waitpid(newPid, &status, 0);

    // Fix up registers for both processes...
    regs.ebx = ebx;
    regs.ecx = ecx;
    regs.edx = edx;
    regs.esi = esi;
    regs.eip = eip;
    ptrace(PTRACE_SETREGS, mPid, NULL, &regs);

    regs.eax = 0;
    ptrace(PTRACE_SETREGS, newPid, NULL, &regs);
    ptrace(PTRACE_SYSCALL, newPid, NULL, NULL);

    processTable[newPid] = new Process(newPid, true);
  }

  void
  traceSyscall()
  {
    user_regs_struct regs;
    ptrace(PTRACE_GETREGS, mPid, NULL, &regs);

    if (mForkHandlingState > FHS_NOT_HANDLING)
    {
      doForkHandling(regs);
      return;
    }

    if (!mInSyscall)
    {
      if (regs.orig_eax == __NR_vfork || regs.orig_eax == __NR_fork)
      {
        startFork(regs);
      }
      return;
    }

    if (regs.orig_eax == __NR_mmap)
    {
      struct mmap_arg_struct args;
      copy_from_user(&args, regs.ebx, sizeof(args));
      traceMmap2(args.addr, args.len, args.prot, args.flags, args.fd,
                 args.offset, regs.eax);
    }
    else if (regs.orig_eax == __NR_mmap2)
    {
      traceMmap2(static_cast<unsigned long>(regs.ebx),
                 static_cast<unsigned long>(regs.ecx),
                 static_cast<unsigned long>(regs.edx),
                 static_cast<unsigned long>(regs.esi),
                 static_cast<unsigned long>(regs.edi),
                 static_cast<unsigned long>(regs.ebp),
                 regs.eax);
    }
    else if (regs.orig_eax == __NR_mremap)
    {
      traceMremap(static_cast<unsigned long>(regs.ebx),
                  static_cast<unsigned long>(regs.ecx),
                  static_cast<unsigned long>(regs.edx),
                  static_cast<unsigned long>(regs.esi),
                  static_cast<unsigned long>(regs.edi),
                  regs.eax);
    }
    else if (regs.orig_eax == __NR_munmap)
    {
      traceMunmap(static_cast<unsigned long>(regs.ebx),
                  static_cast<unsigned long>(regs.ecx),
                  regs.eax);
    }
    else if (regs.orig_eax == __NR_brk)
    {
      traceBrk(static_cast<unsigned long>(regs.ebx));
    }
  }

  void
  copy_from_user(void* here, long there, size_t length)
  {
    if (length & 0x3 != 0)
    {
      length = (length | 0x3) + 1;
    }

    unsigned long* p = reinterpret_cast<unsigned long*>(here);

    size_t i;
    for (i = 0; i < length; i += 4)
      *p++ = ptrace(PTRACE_PEEKDATA, mPid, there + i, NULL);
  }

  void
  copy_to_user(long there, const void* here, size_t length)
  {
    if (length & 0x3 != 0)
    {
      length = (length | 0x3) + 1;
    }

    const unsigned long* p = reinterpret_cast<const unsigned long*>(here);

    size_t i;
    for (i = 0; i < length; i += 4)
      ptrace(PTRACE_POKEDATA, mPid, there + i, *p++);
  }

  void
  changeProcessAllocation(int change)
  {
    ChangeGlobalAllocation(change);
    mAllocatedBytes += change;
  }

  void
  traceMmap2
  (
   unsigned long addr, unsigned long len,
   unsigned long prot, unsigned long flags,
   unsigned long fd, unsigned long pgoff,
   unsigned long actual_addr
  )
  {
    changeProcessAllocation(len);
  }

  void
  traceMremap
  (
   unsigned long addr, unsigned long old_len,
   unsigned long new_len, unsigned long flags,
   unsigned long new_addr,
   long actual_addr
  )
  {
    changeProcessAllocation(new_len - old_len);
  }

  void
  traceMunmap(unsigned long addr, unsigned long len,
              long retval)
  {
    changeProcessAllocation(-len);
  }

  void
  startFork(struct user_regs_struct& regs)
  {
    // Save state because it will get mangled...
    eax = regs.eax;
    ebx = regs.ebx;
    ecx = regs.ecx;
    edx = regs.edx;
    esi = regs.esi;
    eip = regs.eip;

    // See if we already allocated a block for injected data...
    if (mForkHandlingState == FHS_NEVER_HANDLED)
    {
      mForkHandlingState = FHS_WAITING_MMAP2;
      regs.orig_eax = __NR_mmap2;
      regs.ebx = 0; // no preference for location
      regs.ecx = 4096; // allocate one page.
      // give it read/write/execute permissions.
      regs.edx = PROT_EXEC | PROT_READ | PROT_WRITE;
      // page is private and anonymous
      regs.esi = MAP_PRIVATE | MAP_ANONYMOUS;
      regs.eip = mForkHandlerAddress;

      // Save registers...
      ptrace(PTRACE_SETREGS, mPid, NULL, &regs);

      // We are going into the kernel for syscall after we return...
      return;
    }
    else
    {
      mForkHandlingState = FHS_WAITING_NOOP;
      regs.orig_eax = 20; // getpid, used as a no-op because we don't want a syscall
                          // any more.
      regs.eip = mForkHandlerAddress;

      // Save registers...
      ptrace(PTRACE_SETREGS, mPid, NULL, &regs);

      // We are going into the kernel for syscall after we return...
      return;
    }
  }

  void
  traceBrk(unsigned long addr)
  {
    if (addr == 0)
      return;
    changeProcessAllocation(addr - mLastBreak);
    mLastBreak = addr;
  }
};

int
main(int argc, char** argv)
{
  char** newArgv = new char*[argc];
  int i;

  if (argc < 2)
  {
    printf("Usage: %s program args...\n",
           argv[0]);
    return 1;
  }

  for (i = 1; i < argc; i++)
    newArgv[i - 1] = argv[i];
  newArgv[argc - 1] = NULL;

  pid_t pid = fork();

  if (pid == 0)
  {
    ptrace(PTRACE_TRACEME, 0, NULL, NULL);
    execv(newArgv[0], newArgv);
    perror("execv");
    return 1;
  }
  if (pid == -1)
  {
    perror("fork");
    return 1;
  }
  
  int status;
  waitpid(pid, &status, 0);
  if (!WIFSTOPPED(status))
  {
    fprintf(stderr, "Process exited early (%u)\n", status);
    fflush(stderr);
    return 2;
  }

  memset(processTable, 0, sizeof(processTable));
  processTable[pid] = new Process(pid);
  processTable[pid]->isMainProcess();

  while (true)
  {
    pid = waitpid(0, &status, 0);
    if (processTable[pid] != 0)
      processTable[pid]->waitReturned(status);
  }
}
